package com.jd.jobkeeper.client;

import com.jd.jobkeeper.protocals.*;
import com.jd.jobkeeper.transport.DefaultTransport;
import com.jd.jobkeeper.transport.HandlerRegistry;
import com.jd.jobkeeper.transport.Transport;
import com.jd.jobkeeper.utils.net.NetUtils;
import io.rsocket.Payload;
import io.rsocket.util.DefaultPayload;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;

import java.lang.management.ManagementFactory;
import java.nio.charset.StandardCharsets;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public class JobWorkerNode {

    private static Logger log = LoggerFactory.getLogger(JobWorkerNode.class);

    private AtomicBoolean started = new AtomicBoolean(false);

    private final HandlerRegistry registry;
    final String nodeId;
    final String ip;
    final Sinks.Many<Task> finishedTask;
    final Scheduler scheduler;
    final JobClientProtocol protocol;
    final JobConfig config;
    final AddressIndex addressIndex;


    public JobWorkerNode(JobConfig config) {
        this.config = config;
        this.ip = NetUtils.findFirstNonLoopbackAddress(true, null).getHostAddress();
        this.nodeId = ip + "$" + UUID.randomUUID().toString();

        this.scheduler = Schedulers.newBoundedElastic(8, 1024, "ss", 60);
        this.addressIndex = new AddressIndex(config, scheduler);
        this.finishedTask = Sinks.many().multicast().onBackpressureBuffer();

        this.protocol = new JobClientProtocolImpl(scheduler, finishedTask);
        this.registry = new HandlerRegistry();
        this.registry.register(protocol, JobClientProtocol.class);
    }


    public synchronized void registry(Handler handler) {
        this.protocol.register(handler);
    }


    public void start() {
        if (started.compareAndSet(false, true)) {
            this.addressIndex.start()
                    .doOnError(t -> {
                        log.info("Get address error", t);
                    })
                    .doOnSuccess(v -> connectAndSubscribe())
                    .block();
        }
    }

    public void stop() {
        if (started.compareAndSet(true, false)) {
            this.addressIndex.stop();
            this.protocol.stop();

        }
    }


    public void connectAndSubscribe() {
        if (!started.get()) {
            return;
        }
        Node node = addressIndex.selectNode();

        buildTransport(node)
                .flatMap(t -> _subscribe(t))
                .subscribe(t -> {
                    protocol.onConnected(t);
                    //this.transport = t;
                    t.onClose().doFinally(s -> {
                        log.info("The transport is close 1....");
                        delayConnectAndSubscribe();
                    }).subscribe();
                }, e -> {
                    log.warn("Fail subscribe to " + node.toString(), e);
                    delayConnectAndSubscribe();
                });
    }

    private void delayConnectAndSubscribe() {
        this.scheduler.schedule(() -> connectAndSubscribe(), 5, TimeUnit.SECONDS);
    }

    private Mono<Transport> _subscribe(Transport t) {
        Request request = Request.newBuilder().setNamespace(this.config.getNamespace()).build();

        return t.send("subscribe", request, SubscribeResponse.class)
                .flatMap(r -> {
                    if (r.getStatus() == ResponseStatus.OK) {
                        return Mono.just(t);
                    } else {
                        t.dispose();
                        Node redirectNode = new Node(r.getServer().getHost(), r.getServer().getPort());
                        log.info("Redirect to node {}", redirectNode);
                        return buildTransport(redirectNode).flatMap(this::_subscribe);
                    }
                });
    }

    private Mono<Transport> buildTransport(Node node) {
        log.info("Connecting to node " + node);
        SetupRequest setupRequest = SetupRequest.newBuilder()
                .setAppName(this.config.getApp())
                .setNamespace(this.config.getNamespace())
                .setId(nodeId)
                .setHost(ip)
                .setPort(18000)
                .build();

        Payload setupPayload = DefaultPayload.create(setupRequest.toByteArray(),
                StandardCharsets.UTF_8.encode("setup").array());

        DefaultTransport transport = new DefaultTransport(node.getHost(), node.getPort());
        transport.setHandlerRegistry(registry);

        return transport.connectWithPayLoad(node.getHost(), node.getPort(), setupPayload)
                .map(rSocket -> new DefaultTransport(rSocket))
                .doOnNext(t -> log.info("Successfully connected to node " + node))
                .cast(Transport.class);
    }

    private String getPid() {
        try {
            String jvmName = ManagementFactory.getRuntimeMXBean().getName();
            return jvmName.split("@")[0];
        } catch (Throwable ex) {
            return null;
        }
    }


    public static void main(String[] args) throws Exception {
        JobConfig jobConfig = new JobConfig();
        jobConfig.setIndex("");
        JobWorkerNode node = new JobWorkerNode(jobConfig);

        node.start();

        synchronized (node) {
            node.wait();
        }
    }
}
